#开发环境：VSCode
#Python 版本 ： 3.7.13（通过Anaconda 配置)
#opencv 版本 ： 4.5.5

import cv2
import numpy as np
import kociemba                              # 请先安装 kociemba，https://github.com/muodov/kociemba
from colordetection import color_detector    # 该文件来源于 https://github.com/kkoomen/qbr

CUBE_HEIGHT = 360

capture = cv2.VideoCapture(0)    #调整参数实现读取视频或调用摄像头

area_frant = [(315, 178),(453, 114),(436, 275),(314, 364)]
area_left = [(171, 112),(315, 178),(314, 364),(184, 283)]
area_top = [(171, 112),(311, 62),(453, 114),(315, 178)]
selected_point = ()
move_to_point = ()

"""
    功能：替换移动的顶点并返回
"""
def get_area_points(point_list):
    result_points = []
    if len(selected_point) == 0 or len(move_to_point) == 0:
        return point_list
    for index in range(len(point_list)):
        if selected_point == point_list[index]:
            result_points.append(move_to_point)
        else:
            result_points.append(point_list[index])
    return result_points

"""
    功能：取得点击坐标
"""
def handle_mouse_event(event,x,y,flags,param):
    global selected_point
    global move_to_point
    global area_frant
    global area_left
    global area_top
    if event == cv2.EVENT_LBUTTONDOWN :
        pts = np.array([[x-3, y-3],[x-3, y+3],[x+3, y+3],[x+3, y-3]], np.int32)

        min_distance = 100
        for point in area_frant + area_left + area_top:
            dist=cv2.pointPolygonTest(pts,point,True)
            if dist >=0 and dist < min_distance:
                # 在顶点3个像素以内，选中该顶点
                min_distance = dist
                selected_point = point
                print(point)
        if min_distance == 100:
            selected_point = ()
    elif event == cv2.EVENT_LBUTTONUP :
        if len(selected_point) > 0 and len(move_to_point) > 0:
            area_frant = get_area_points(area_frant)
            area_left = get_area_points(area_left)
            area_top = get_area_points(area_top)
            selected_point = move_to_point
        move_to_point = ()
    elif event == cv2.EVENT_MOUSEMOVE and flags & cv2.EVENT_FLAG_LBUTTON:
        if len(selected_point) > 0:
            move_to_point = (x,y)
    return

cv2.namedWindow("Cube detect")
cv2.setMouseCallback("Cube detect",handle_mouse_event)

"""
    功能：识别魔方面上的九个色块颜色
"""
def detect_colors(frame,contours):
    result_colors = []
    for contour in contours:
        (x,y,w,h) = contour
        roi = frame[y+30:y+h-30, x+30:x+w-30]
        avg_bgr = color_detector.get_dominant_color(roi)
        closest_color = color_detector.get_closest_color(avg_bgr)['color_bgr']
        result_colors.append(closest_color)
    return result_colors

"""
    功能：找到魔方某一面的九个方块坐标
"""
def find_contours(image1):
    cv2.line(image1, (int(CUBE_HEIGHT/3)-3,0),(int(CUBE_HEIGHT/3)-3,CUBE_HEIGHT), (0,0,0),5) 
    cv2.line(image1, (int(CUBE_HEIGHT*2/3)-3,0),(int(CUBE_HEIGHT*2/3)-3,CUBE_HEIGHT), (0,0,0),5) 
    cv2.line(image1, (0,int(CUBE_HEIGHT/3)-3),(CUBE_HEIGHT,int(CUBE_HEIGHT/3)-3), (0,0,0),5) 
    cv2.line(image1, (0,int(CUBE_HEIGHT*2/3)-3),(CUBE_HEIGHT,int(CUBE_HEIGHT*2/3)-3), (0,0,0),5) 
    cv2.rectangle(image1,(0,0),(CUBE_HEIGHT,CUBE_HEIGHT),(0,0,0),10)

    grayFrame = cv2.cvtColor(image1, cv2.COLOR_BGR2GRAY)
    blurredFrame = cv2.blur(grayFrame, (3, 3))
    cannyFrame = cv2.Canny(blurredFrame, 30, 60, 3)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (9, 9))
    dilatedFrame = cv2.dilate(cannyFrame, kernel)

    """Find the contours of a 3x3x3 cube."""
    contours, hierarchy = cv2.findContours(dilatedFrame, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    final_contours = []
    # Step 1/2: 找到9个矩形.
    for contour in contours:
        perimeter = cv2.arcLength(contour, True) # 计算轮廓周长
        approx = cv2.approxPolyDP(contour, 0.1 * perimeter, True)  # 多边形拟合，拟合精度为周长的1/10
        if len (approx) == 4 and perimeter > CUBE_HEIGHT and perimeter < CUBE_HEIGHT*2:
            # cv2.polylines(image,[approx],True,(255,0,0),3)
            area = cv2.contourArea(contour)  # 计算面积
            (x, y, w, h) = cv2.boundingRect(approx) # 寻找包裹轮廓的最小正矩形

            # Find aspect ratio of boundary rectangle around the countours.
            ratio = w / float(h)  # 计算长宽比

            # Check if contour is close to a square.
            # 长宽比范围在 0.8 至 1.2 之间，认为是矩形
            # 宽度在 90至 120 之间则矩形尺寸比较合适
            # 多边形面积与矩形面积相似
            if ratio >= 0.8 and ratio <= 1.2 and w >= CUBE_HEIGHT/3-30 and w <= CUBE_HEIGHT/3 and area / (w * h) > 0.4:
                final_contours.append((x, y, w, h))

    # Return early if we didn't found 9 or more contours.
    if len(final_contours) < 9:
        return []

    # 将找到的9个矩形按左上至右下进行排序返回
    # Sort contours on the y-value first.
    y_sorted = sorted(final_contours, key=lambda item: item[1])

    # Split into 3 rows and sort each row on the x-value.
    top_row = sorted(y_sorted[0:3], key=lambda item: item[0])
    middle_row = sorted(y_sorted[3:6], key=lambda item: item[0])
    bottom_row = sorted(y_sorted[6:9], key=lambda item: item[0])

    sorted_contours = top_row + middle_row + bottom_row
    return sorted_contours

"""
    功能：取得魔方指定面
"""
def detect_cube_side(frame,start_pts):
    # 旋转透视
    finally_pts = np.float32([[0,0],[CUBE_HEIGHT,0],[CUBE_HEIGHT,CUBE_HEIGHT],[0,CUBE_HEIGHT]]);
    matrix = cv2.getPerspectiveTransform(start_pts,finally_pts)
    image = cv2.warpPerspective(frame,matrix,(CUBE_HEIGHT,CUBE_HEIGHT))

    contours = find_contours(image)
    if len(contours) == 9:
        # 生成预览图
        preview = np.copy(image)

        contour_colors = detect_colors(image, contours)
        # 画出9个方块
        for index,(x,y,w,h) in enumerate(contours): 
            cv2.rectangle(preview,(x,y),(x+w,y+h),contour_colors[index],-1)

        return (image, color_detector.get_closest_color(contour_colors[4])['color_name'], contours, preview)
    return (image,0,0,image)

"""
    功能：提高画面亮度
"""
def light_color_pallete(image,gamma = 0):
    # 通过gamma变换调整画面亮度，提高识别准确率
    lookUpTable = np.empty((1,256), np.uint8)
    for i in range(256):
        lookUpTable[0,i] = np.clip(pow(i / 255.0, gamma) * 255.0, 0, 255)
    return cv2.LUT(image, lookUpTable)

""""
    功能：根据检测的6个面中心块颜色校准色板
"""
def reset_color_pallete(cube_sides):
    calibrated_colors = {}
    for side in ['white', 'red', 'green', 'yellow', 'orange', 'blue']:
        image = cube_sides[side]
        contours = cube_sides[side+"_contours"]
        (x,y,w,h) = contours[4] # 中心块
        roi = image[y+7:y+h-7, x+14:x+w-14]
        avg_bgr = color_detector.get_dominant_color(roi)
        calibrated_colors[side] = avg_bgr
    color_detector.set_cube_color_pallete(calibrated_colors)

"""
    功能：获取魔方还原步骤
"""
def solve_cube(cube_sides):
    cubeState = ""
    for side in ['white', 'red', 'green', 'yellow', 'orange', 'blue']:
        # if side+"_preview" not in cube_sides:
        #     return
        # image1 = cube_sides[side+"_preview"]
        image = cube_sides[side]
        contour_colors = detect_colors(image, cube_sides[side+"_contours"]) # 获取9个块的颜色
        for index in range(0,9):
            cubeState += ''.join(color_detector.convert_bgr_to_notation(contour_colors[index]))
    print(cubeState)

    try:
        print(kociemba.solve(cubeState))
        return kociemba.solve(cubeState)
    except Exception:
        print("Solve Failed！")
        return ""

"""
    主程序
"""
if __name__ == '__main__':
    cube_sides = {}
    detect_step = 0
    preview_Y = 370
    preview_X = 430
    gamma = 0.7
    solution = ""
    while 1:
        ret, frame = capture.read()
        key = cv2.waitKey(100) & 0xff
        if  key == 27:  # ESC
            break
        elif key == 32 and detect_step > 0: # 空格
            detect_step = detect_step -1

        if gamma < 1:
            frame = light_color_pallete(frame,gamma)
        front_pts = np.float32(area_frant)
        if detect_step > 0:
            front_pts = np.float32([area_frant[2],area_frant[3],area_frant[0],area_frant[1]])
        (imageFront, colorFront, contoursFront, previewFront) = detect_cube_side(frame,front_pts)

        left_pts = np.float32([area_left])
        if detect_step > 0:
            left_pts = np.float32([area_left[2],area_left[3],area_left[0],area_left[1]])
        (imageLeft, colorLeft, contoursLeft, previewLeft) = detect_cube_side(frame,left_pts)

        top_pts = np.float32(area_top)
        if detect_step > 0:
            top_pts = np.float32([area_top[1],area_top[2],area_top[3],area_top[0]])
        (imageTop, colorTop,contoursTop, previewTop) = detect_cube_side(frame,top_pts)

        if detect_step == 0 and colorFront == "green" and colorLeft == "red" and colorTop == "white":
            # 有时橙色块会被识别为红色，需要将其矫正过来
            new_palette = color_detector.cube_color_palette
            (x,y,w,h) = contoursLeft[4] # 中心块
            roi = imageLeft[y+7:y+h-7, x+14:x+w-14]
            avg_bgr = color_detector.get_dominant_color(roi)
            new_palette['orange'] = avg_bgr
            print(new_palette)
            color_detector.set_cube_color_pallete(new_palette)
        if detect_step ==0 and colorFront == "green" and colorLeft == "orange" and colorTop == "white":
            # 绿 橙 白 三面检测完毕
            cube_sides["green"] = imageFront
            cube_sides["green_contours"] = contoursFront
            cube_sides["green_preview"] = previewFront
            cube_sides["orange"] = imageLeft
            cube_sides["orange_contours"] = contoursLeft
            cube_sides["orange_preview"] = previewLeft
            cube_sides["white"] = imageTop
            cube_sides["white_contours"] = contoursTop
            cube_sides["white_preview"] = previewTop
            detect_step = 1
        elif detect_step == 1 and colorFront == "red" and colorLeft == "blue" and colorTop == "yellow":
            # 红 蓝 黄 三面检测完毕
            cube_sides["red"] = imageFront
            cube_sides["red_contours"] = contoursFront
            cube_sides["red_preview"] = previewFront
            cube_sides["blue"] = imageLeft
            cube_sides["blue_contours"] = contoursLeft
            cube_sides["blue_preview"] = previewLeft
            cube_sides["yellow"] = imageTop
            cube_sides["yellow_contours"] = contoursTop
            cube_sides["yellow_preview"] = previewTop
            reset_color_pallete(cube_sides)
            solution = solve_cube(cube_sides)
            if len(solution) > 0 :
                detect_step = 2
            

        # 画正面多边形
        points=np.array(get_area_points(area_frant),np.int32)  #多边形的顶点坐标
        cv2.polylines(frame,[points],True,(0,255,0),2)

        # 画左面多边形
        points=np.array(get_area_points(area_left),np.int32)  #多边形的顶点坐标
        cv2.polylines(frame,[points],True,(0,255,0),2)  

        # 画上面多边形
        points=np.array(get_area_points(area_top),np.int32)  #多边形的顶点坐标
        cv2.polylines(frame,[points],True,(0,255,0),2)  

        # 画出选中的点
        if len(selected_point) > 0:
            if len(move_to_point) > 0:
                cv2.circle(frame,move_to_point,3,(0,0,255),2)
            else:
                cv2.circle(frame,selected_point,3,(0,0,255),2)

        tips = ["1.Change vertex by mouse and put cube on assigned position.",
                "2.Flip cube and put on assigned position.",
                "  Cube detect SUCCESS!"]
        cv2.putText(frame,tips[detect_step],(0, 30),0,color=(255, 0, 0),fontScale=0.6,thickness=2)
        if detect_step == 2:
            cv2.putText(frame,solution,(20, 60),0,color=(255, 0, 0),fontScale=0.6,thickness=2)
        cv2.putText(frame,"Press SPACE to last step.",(20, 450),0,color=(255, 255, 255),fontScale=0.6,thickness=2)
        cv2.putText(frame,"Press ESC to quit.",(20, 470),0,color=(255, 255, 255),fontScale=0.6,thickness=2)
        
        M = cv2.moments(front_pts)
        frant_center = (int(M['m10']/M['m00']), int(M['m01']/M['m00']))
        M = cv2.moments(left_pts)
        left_center = (int(M['m10']/M['m00']), int(M['m01']/M['m00']))
        M = cv2.moments(top_pts)
        top_center = (int(M['m10']/M['m00']), int(M['m01']/M['m00']))
        color_palette = {
            'red'   : (0, 0, 255),
            'orange': (0, 165, 255),
            'blue'  : (255, 0, 0),
            'green' : (0, 255, 0),
            'white' : (255, 255, 255),
            'yellow': (0, 255, 255)
        }
        if detect_step == 0 :
            cv2.circle(frame,frant_center,10,color_palette['green'],-1)
            cv2.circle(frame,left_center,10,color_palette['orange'],-1)
            cv2.circle(frame,top_center,10,color_palette['white'],-1)
            if len(selected_point) == 0:
                for point in area_frant + area_left +area_top:
                    cv2.circle(frame,point,3,color_palette['blue'],2)
        else:
            cv2.circle(frame,frant_center,10,color_palette['red'],-1)
            cv2.circle(frame,left_center,10,color_palette['blue'],-1)
            cv2.circle(frame,top_center,10,color_palette['yellow'],-1)

        
        # 在右下角显示已经识别出来的各个面
        if "orange_preview" in cube_sides:
            frame[preview_Y:preview_Y+50,preview_X:preview_X+50,:3] = cv2.resize(cube_sides["orange"],(50,50))
            frame[preview_Y:preview_Y+50,preview_X+52:preview_X+102,:3] = cv2.resize(cube_sides["green"],(50,50))
            frame[preview_Y-52:preview_Y-2,preview_X+52:preview_X+102,:3] = cv2.resize(cube_sides["white"],(50,50))
        else:
            cv2.rectangle(frame,(preview_X,preview_Y),(preview_X + 50,preview_Y+50),(120,120,120),1)
            cv2.circle(frame,(preview_X+25,preview_Y+25),10,color_palette['orange'],-1)

            cv2.rectangle(frame,(preview_X + 52,preview_Y),(preview_X + 102,preview_Y+50),(120,120,120),1)
            cv2.circle(frame,(preview_X + 52+25,preview_Y+25),10,color_palette['green'],-1)

            cv2.rectangle(frame,(preview_X + 52,preview_Y - 52),(preview_X +102,preview_Y-2),(120,120,120),1)
            cv2.circle(frame,(preview_X + 52+25,preview_Y - 52+25),10,color_palette['white'],-1)
        if "yellow_preview" in cube_sides:
            frame[preview_Y+52:preview_Y+102,preview_X+52:preview_X+102,:3] = cv2.resize(cube_sides["yellow"],(50,50))
            frame[preview_Y:preview_Y+50,preview_X+104:preview_X+154,:3] = cv2.resize(cube_sides["red"],(50,50))
            frame[preview_Y:preview_Y+50,preview_X+156:preview_X+206,:3] = cv2.resize(cube_sides["blue"],(50,50))
        else:
            cv2.rectangle(frame,(preview_X + 52,preview_Y + 52),(preview_X + 102,preview_Y+102),(120,120,120),1)
            cv2.circle(frame,(preview_X + 52+25,preview_Y + 52+25),10,color_palette['yellow'],-1)

            cv2.rectangle(frame,(preview_X + 104,preview_Y),(preview_X + 154,preview_Y+50),(120,120,120),1)
            cv2.circle(frame,(preview_X + 104+25,preview_Y + 25),10,color_palette['red'],-1)

            cv2.rectangle(frame,(preview_X + 156,preview_Y),(preview_X + 206,preview_Y+50),(120,120,120),1)
            cv2.circle(frame,(preview_X + 156+25,preview_Y + 25),10,color_palette['blue'],-1)

        cv2.imshow("Cube detect", frame)

    capture.release()
    cv2.destroyAllWindows()